package com.koicarp.agent.rag;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

import cn.hutool.core.collection.ListUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.tool.ToolExecution;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.MutationResult;
import io.milvus.param.R;
import io.milvus.param.dml.InsertParam;
import kotlin.contracts.Returns;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

@RestController
@RequestMapping("api")
public class ApiController {

	static String MODEL_NAME = "deepseek-r1:7b";
	
	static String BASE_URL = "http://192.168.0.67:11434";

	@RequestMapping("chat")
	public String  chat(@RequestBody String content) {
		List<Document> documents = FileSystemDocumentLoader.loadDocumentsRecursively("C:\\Users\\dell.pc\\Desktop\\deepseek_doc\\two");
		InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
		EmbeddingStoreIngestor.ingest(documents, embeddingStore);
		
		OllamaChatModel model = OllamaChatModel.builder()
		.baseUrl(BASE_URL)
		.modelName(MODEL_NAME)
		.timeout(Duration.ofHours(10))
		.build();
		Assistant assistant = AiServices.builder(Assistant.class)
				.chatLanguageModel(model)
//				.chatMemory(MessageWindowChatMemory.withMaxMessages(10))
				.contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
			    .build();
		
		 String token = assistant.chat(content);
		return token;
	}

	@RequestMapping("streamChat")
	public Flux<String> streamChat(@RequestBody String content) {
		List<Document> documents = FileSystemDocumentLoader.loadDocumentsRecursively("C:\\Users\\dell.pc\\Desktop\\deepseek_doc\\two");
		InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
		EmbeddingStoreIngestor.ingest(documents, embeddingStore);
		
		OllamaStreamingChatModel model = OllamaStreamingChatModel.builder()
		.baseUrl(BASE_URL)
		.modelName(MODEL_NAME)
		.timeout(Duration.ofHours(10))
		.build();
		
		StreamAssistant assistant = AiServices.builder(StreamAssistant.class)
		.streamingChatLanguageModel(model)
		.contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
		.build();
		// 创建 StreamingResponseBody
		Sinks.Many<String> sink = Sinks.many().unicast().onBackpressureBuffer();
		TokenStream tokenStream = assistant.chat(content);
		tokenStream.onNext(token->{
			sink.tryEmitNext(token);
		})
		.onComplete(message->{
			sink.tryEmitComplete();
		})
		.onError(t->{
			sink.tryEmitError(t);
		}).start();
		
		return sink.asFlux();
	}

	@RequestMapping("get")
	public Flux<String> getChat() {
		return streamChat("前车间中部搅拌电机电源接口未封堵，在提供的信息中，违反了哪个条款？");
	}
	
	
	@RequestMapping("milvusChat")
	public String milvusChat(@RequestBody String content) {
//		List<Document> documents = FileSystemDocumentLoader.loadDocumentsRecursively("C:\\Users\\dell.pc\\Desktop\\deepseek_doc\\two");
		
		MilvusEmbeddingStore embeddingStore = MilvusEmbeddingStore.builder()
		.uri("http://192.168.0.67")
		.collectionName("trouble_20250221")
		.dimension(384)
		.port(19530)
		.build();
		
		
		AllMiniLmL6V2EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
//		EmbeddingStoreIngestor build = EmbeddingStoreIngestor.builder()
//		.documentSplitter(new TroubleJsonDocumentSplitter())
//		.embeddingStore(embeddingStore)
//		.embeddingModel(embeddingModel)
//		.build();
//		build.ingest(documents);
		
		OllamaChatModel model = OllamaChatModel.builder()
		.baseUrl(BASE_URL)
		.modelName(MODEL_NAME)
		.build();
		
		EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder().embeddingModel(embeddingModel)
//				.maxResults(3)
				.minScore(0.6)
		.embeddingStore(embeddingStore).build();

		
		Assistant assistant = AiServices.builder(Assistant.class)
		.chatLanguageModel(model)
		.contentRetriever(contentRetriever)
		.build();
		
		
		return assistant.chat(content);
	}
	
	
	@RequestMapping("streamMilvusChat")
	public Flux<String> streamMilvusChat(@RequestBody String content) {
//		List<Document> documents = FileSystemDocumentLoader.loadDocumentsRecursively("C:\\Users\\dell.pc\\Desktop\\deepseek_doc\\two");
		
		MilvusEmbeddingStore embeddingStore = MilvusEmbeddingStore.builder()
		.uri("http://192.168.0.67")
		.collectionName("trouble_20250221")
		.dimension(384)
		.port(19530)
		.build();
		
		
		AllMiniLmL6V2EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
//		EmbeddingStoreIngestor build = EmbeddingStoreIngestor.builder()
//		.documentSplitter(new TroubleJsonDocumentSplitter())
//		.embeddingStore(embeddingStore)
//		.embeddingModel(embeddingModel)
//		.build();
//		build.ingest(documents);
		
		OllamaStreamingChatModel model = OllamaStreamingChatModel.builder()
		.baseUrl(BASE_URL)
		.modelName(MODEL_NAME)
		.timeout(Duration.ofHours(10))
		.build();
		
		EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder().embeddingModel(embeddingModel)
//				.maxResults(3)
				.minScore(0.6)
		.embeddingStore(embeddingStore).build();

		
		
		StreamAssistant assistant = AiServices.builder(StreamAssistant.class)
				.streamingChatLanguageModel(model)
				.contentRetriever(contentRetriever)
				.build();
				// 创建 StreamingResponseBody
				Sinks.Many<String> sink = Sinks.many().unicast().onBackpressureBuffer();
				TokenStream tokenStream = assistant.chat(content);
				tokenStream.onNext(token->{
					sink.tryEmitNext(token);
				})
				.onComplete(message->{
					sink.tryEmitComplete();
				})
				.onError(t->{
					sink.tryEmitError(t);
				}).start();
				
				return sink.asFlux();
	}

	@RequestMapping("add")
	public void add() {
		List<Document> documents = FileSystemDocumentLoader.loadDocumentsRecursively("C:\\Users\\dell.pc\\Desktop\\deepseek_doc\\two");
		TroubleJsonDocumentSplitter splitter = new TroubleJsonDocumentSplitter();
		List<TextSegment> segments = splitter.splitAll(documents);

		AllMiniLmL6V2EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
		
		List<TextSegment> troubles = new ArrayList<>();
		for (TextSegment textSegment : segments) {
			String text = textSegment.text();
			JSONObject parseObj = JSONUtil.parseObj(text);
			String desc = parseObj.getStr("隐患描述");
			troubles.add(TextSegment.from(desc));
		}
		List<Embedding> contentsEmbeddings = embeddingModel.embedAll(troubles).content();
		
		
		MilvusEmbeddingStore embeddingStore = MilvusEmbeddingStore.builder()
		.uri("http://192.168.0.67")
		.collectionName("trouble_20250221")
		.dimension(384)
		.port(19530)
		.build();
		
		embeddingStore.addAll(contentsEmbeddings, segments);
	}
}
